import pandas as pd
import numpy as np
import os
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy.stats import zscore


# --- 数据加载和基础处理函数 ---
def load_data_in_range(start_time_str: str, end_time_str: str, base_path: str, symbol: str,
                       timeframe: str) -> pd.DataFrame:
    """从CSV文件中加载指定时间范围内的数据。"""
    start_date = pd.to_datetime(start_time_str)
    end_date = pd.to_datetime(end_time_str)
    data_folder_path = os.path.join(base_path, symbol, timeframe)
    if not os.path.isdir(data_folder_path):
        print(f"错误: 文件夹不存在 -> {data_folder_path}")
        return pd.DataFrame()

    date_range = pd.date_range(start=start_date.date(), end=end_date.date(), freq='D').strftime('%Y-%m').unique()
    if len(date_range) == 0: date_range = [start_date.strftime('%Y-%m')]

    all_dfs = []
    print(f"将在以下月份文件中查找数据: {list(date_range)}")
    for year_month in date_range:
        file_path = os.path.join(data_folder_path, f"{symbol.upper()}-{timeframe}-{year_month}.csv")
        if os.path.exists(file_path):
            print(f"正在读取文件: {file_path}")
            try:
                # 只需要加载DDI计算所必需的列
                df = pd.read_csv(file_path, usecols=['open_time', 'open', 'high', 'low', 'close', 'volume'])
                all_dfs.append(df)
            except Exception as e:
                print(f"读取或处理文件 {file_path} 时出错: {e}")
        else:
            print(f"警告: 未找到文件 {file_path}")

    if not all_dfs:
        print("错误: 在指定时间范围内未找到任何数据文件。")
        return pd.DataFrame()

    combined_df = pd.concat(all_dfs, ignore_index=True)
    combined_df.drop_duplicates(subset=['open_time'], inplace=True)
    return combined_df


def process_and_filter_data(df: pd.DataFrame, start_time_str: str, end_time_str: str) -> pd.DataFrame:
    """对数据进行清洗、格式化和时间过滤。"""
    if df.empty:
        return df
    df['open_time'] = pd.to_datetime(df['open_time'], unit='ms')
    df.set_index('open_time', inplace=True)
    numeric_cols = ['open', 'high', 'low', 'close', 'volume']
    df[numeric_cols] = df[numeric_cols].apply(pd.to_numeric, errors='coerce')
    df.dropna(subset=numeric_cols, inplace=True)
    df.sort_index(inplace=True)
    start_ts = pd.to_datetime(start_time_str)
    end_ts = pd.to_datetime(end_time_str)
    df_filtered = df.loc[start_ts:end_ts]
    print(f"数据处理完成。筛选出 {len(df_filtered)} 条K线。")
    return df_filtered


# --- 新增：DDI 指标计算函数 ---
def calculate_ddi(df: pd.DataFrame, z_period: int = 50, ema_period: int = 5, threshold: float = 1.0) -> pd.DataFrame:
    """
    计算爆裂衰减指数 (Detonation Decay Index, DDI)。
    """
    if df.empty:
        return df

    # --- 1. 计算瞬时爆裂能量 (IDE) ---
    df['normalized_range'] = (df['close'] - df['open']).abs() / df['close']
    # 使用滚动Z-score来避免前视偏差
    df['z_range'] = df['normalized_range'].rolling(window=z_period).apply(
        lambda x: (x[-1] - x.mean()) / x.std() if x.std() > 0 else 0, raw=True)
    df['z_volume'] = df['volume'].rolling(window=z_period).apply(
        lambda x: (x[-1] - x.mean()) / x.std() if x.std() > 0 else 0, raw=True)

    df['ide'] = (df['z_range'] * df['z_volume']).fillna(0)

    # --- 2. 构建爆裂线 (Detonation Line) ---
    direction = np.sign(df['close'] - df['open'])
    df['detonation_line'] = (df['ide'] * direction).ewm(span=ema_period, adjust=False).mean()

    # --- 3. 构建衰减线 (Decay Line) - 使用高效的矢量化方法 ---
    is_silent = df['ide'] < threshold
    # 为每个连续的“沉寂”或“非沉寂”块创建一个唯一的ID
    block_id = (is_silent.diff() != 0).cumsum()
    # 在每个块内进行累积计数，然后将非沉寂期的计数重置为0
    df['decay_line'] = (is_silent.groupby(block_id).cumsum()).where(is_silent, 0)

    print(f"DDI (Z-{z_period}, EMA-{ema_period}, T-{threshold}) 指标计算完成。")
    return df


# --- 新增：绘制K线和DDI的绘图函数 ---
def plot_kline_and_ddi(df: pd.DataFrame, chart_title: str):
    """
    绘制K线图以及下方的DDI指标（爆裂线和衰减线）。
    """
    if df.empty or 'detonation_line' not in df.columns or 'decay_line' not in df.columns:
        print("数据为空或缺少DDI指标，无法绘制图表。")
        return

    # 创建图表，3行1列，共享X轴
    fig = make_subplots(
        rows=3, cols=1,
        shared_xaxes=True,
        vertical_spacing=0.03,
        row_heights=[0.6, 0.2, 0.2]  # K线图占60%，两个指标各占20%
    )

    # 1. K线图
    fig.add_trace(go.Candlestick(x=df.index,
                                 open=df['open'], high=df['high'],
                                 low=df['low'], close=df['close'],
                                 name='OHLC'),
                  row=1, col=1)

    # 2. 爆裂线 (Detonation Line)
    detonation_colors = ['green' if v >= 0 else 'red' for v in df['detonation_line']]
    fig.add_trace(go.Bar(x=df.index, y=df['detonation_line'],
                         marker_color=detonation_colors,
                         name='Detonation Line'),
                  row=2, col=1)
    fig.add_hline(y=0, line_dash="dash", line_color="grey", row=2, col=1)

    # 3. 衰减线 (Decay Line)
    fig.add_trace(go.Scatter(x=df.index, y=df['decay_line'],
                              fill='tozeroy',  # 创建面积图效果
                              line=dict(color='deepskyblue'),
                              name='Decay Line'),
                  row=3, col=1)

    # 更新布局
    start_str = df.index[0].strftime('%Y-%m-%d %H:%M')
    end_str = df.index[-1].strftime('%Y-%m-%d %H:%M')
    full_title = f'{chart_title} ({start_str} to {end_str}) with Detonation Decay Index (DDI)'

    fig.update_layout(
        title_text=full_title,
        xaxis_rangeslider_visible=False,
        showlegend=False,
        yaxis1_title='Price',
        yaxis2_title='Detonation',
        yaxis3_title='Decay',
        hovermode='x unified',
        template='plotly_dark'
    )
    fig.update_xaxes(title_text='Time', row=3, col=1)

    fig.show()


if __name__ == '__main__':
    # ==================== 用户配置区 ====================
    BASE_DATA_PATH = "F:/personal/binance_klines"
    SYMBOL = "ETHUSDT"
    TIMEFRAME = "30m" # 建议用稍大周期观察结构
    START_TIME = "2025-06-02 00:00:00"
    END_TIME = "2025-08-03 00:00:00"

    # DDI 指标参数，可自行调整
    ZSCORE_PERIOD = 50      # Z-score计算周期，越长对“异常”的定义越苛刻
    EMA_PERIOD = 5          # 爆裂线的平滑周期，越短反应越灵敏
    IDE_THRESHOLD = 1.0     # 定义“爆裂事件”的阈值，可根据历史数据分布调整
    # ====================================================

    print("--- DDI 指标分析程序启动 ---")

    print(f"\n--- 正在加载图表数据 ({START_TIME} to {END_TIME}) ---")
    raw_data_df = load_data_in_range(START_TIME, END_TIME, BASE_DATA_PATH, SYMBOL, TIMEFRAME)
    final_df = process_and_filter_data(raw_data_df, START_TIME, END_TIME)

    if not final_df.empty:
        # 计算DDI指标
        final_df = calculate_ddi(final_df,
                                 z_period=ZSCORE_PERIOD,
                                 ema_period=EMA_PERIOD,
                                 threshold=IDE_THRESHOLD)

        # 绘图
        chart_title_str = f"{SYMBOL} - {TIMEFRAME}"
        plot_kline_and_ddi(final_df, chart_title=chart_title_str)
        print("\n--- 图表生成完毕 ---")
    else:
        print("\n未能加载或处理任何图表数据，程序退出。请检查您的配置和文件路径。")